{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ct3qhBX7rKDC"
      },
      "source": [
        "Copyright 2025 Google LLC."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "z11nEz9krMGQ"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pkOtOq0jDE0c"
      },
      "source": [
        "<table align=\"left\">\n",
        "  <td>\n",
        "      <a target=\"_blank\" href=\"https://colab.research.google.com/github/google-gemini/gemma-cookbook/blob/main/Demos/Emoji-Gemma-on-Web/resources/Convert_Gemma_3_270M_to_LiteRT_for_MediaPipe_LLM_Inference_API.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wDkoc9st4TYk"
      },
      "source": [
        "# Convert Gemma 3 270M to LiteRT for use with MediaPipe LLM Inference API\n",
        "\n",
        "This notebook converts a Gemma 3 270M for use with the [MediaPipe LLM Inference API](https://ai.google.dev/edge/mediapipe/solutions/genai/llm_inference), a library that enables inference on mobile devices or in web browsers. The entire process takes about 15 minutes:\n",
        "\n",
        "1. Set up the Colab environment\n",
        "2. Load the model from Hugging Face\n",
        "3. Convert the model with the AI Edge Torch converter\n",
        "4. Package the model with the MediaPipe Task bundler\n",
        "5. Download the model\n",
        "\n",
        "Gemma 3 270M is designed for task-specific fine-tuning and engineered for efficient performance on mobile, web, and edge devices. You can fine-tune your own model using this [notebook](https://colab.research.google.com/github/google-gemini/gemma-cookbook/blob/main/Demos/Emoji-Gemma-on-Web/resources/Fine_tune_Gemma_3_270M_for_emoji_generation.ipynb) and run it in a demo [web app](https://github.com/google-gemini/gemma-cookbook/tree/main/Demos/Emoji-Gemma-on-Web/app-mediapipe) once converted.\n",
        "\n",
        "## Set up development environment\n",
        "\n",
        "The first step is to install packages using pip."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oEAwOswpsqeq"
      },
      "outputs": [],
      "source": [
        "%pip uninstall -y tensorflow\n",
        "%pip install -U tf-nightly==2.21.0.dev20250819 ai-edge-torch==0.6.0 protobuf transformers\n",
        "%pip install -U jax jaxlib"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q8U1zoAY-4g4"
      },
      "source": [
        "Restart the session runtime to ensure you're using the newly installed packages."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KtibyxZmZjHG"
      },
      "source": [
        "## Load the model\n",
        "To access models on Hugging Face, provide your [Access Token](https://huggingface.co/settings/tokens). You can store it as a Colab secret in the left toolbar by specifying `HF_TOKEN` as the 'Name' and adding your unique token as the 'Value'."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sCQO6oy41gZs"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from google.colab import userdata\n",
        "from huggingface_hub import login\n",
        "hf_token = userdata.get('HF_TOKEN')\n",
        "login(hf_token)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gk2ni3i11LRY"
      },
      "source": [
        "Specify the Hugging Face repo ID of the model to convert. It'll be saved to your Colab files for conversion."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u4p7yHYZ1BKh"
      },
      "outputs": [],
      "source": [
        "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
        "\n",
        "model_author = \"\"                                         #@param {type:\"string\"}\n",
        "model_name = \"myemoji-gemma-3-270m-it\"                    #@param {type:\"string\"}\n",
        "\n",
        "repo_id = f\"{model_author}/{model_name}\"                  # Model to convert\n",
        "save_path = f\"/content/{model_name}\"                      # Path to save resized model\n",
        "\n",
        "model = AutoModelForCausalLM.from_pretrained(repo_id)     # Load the model\n",
        "tokenizer = AutoTokenizer.from_pretrained(repo_id)        # Load the tokenizer\n",
        "\n",
        "model.save_pretrained(save_path)\n",
        "tokenizer.save_pretrained(save_path)\n",
        "\n",
        "print(f\"Model and tokenizer saved to {save_path}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zTvHIgTT-KzZ"
      },
      "source": [
        "## Convert the model\n",
        "Convert and quantize the model using the [AI Edge Torch](https://github.com/google-ai-edge/ai-edge-torch) converter. You can adjust the conversions parameters based on your task's requirements:\n",
        "\n",
        "* `prefill_seq_len`: maximum length of supported input\n",
        "* `kv_cache_max_len`: maximum of prefill + decode context length\n",
        "* `quantize`: the quantization scheme. 8-bit integer quantization (INT8) is good for web environments\n",
        "\n",
        "This takes about 10 minutes. The .tflite model will be saved temporarily to your Colab files."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_vzTdvia1n7E"
      },
      "outputs": [],
      "source": [
        "from ai_edge_torch.generative.examples.gemma3 import gemma3\n",
        "from ai_edge_torch.generative.utilities import converter\n",
        "from ai_edge_torch.generative.utilities.export_config import ExportConfig\n",
        "from ai_edge_torch.generative.layers import kv_cache\n",
        "\n",
        "# Get model, set export settings, and convert to .tflite\n",
        "pytorch_model = gemma3.build_model_270m(save_path)\n",
        "export_config = ExportConfig()\n",
        "export_config.kvcache_layout = kv_cache.KV_LAYOUT_TRANSPOSED\n",
        "export_config.mask_as_input = True\n",
        "converter.convert_to_tflite(\n",
        "    pytorch_model,\n",
        "    output_path=\"/content\",\n",
        "    output_name_prefix=model_name,\n",
        "    prefill_seq_len=128,\n",
        "    kv_cache_max_len=512,\n",
        "    quantize=\"dynamic_int8\",\n",
        "    export_config=export_config,\n",
        ")\n",
        "\n",
        "print (f\"Model converted to .tflite and saved to {save_path}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mqFSmQxi-rCF"
      },
      "source": [
        "## Create a MediaPipe Task Bundle\n",
        "\n",
        "A MediaPipe Task file (.task) bundles the original model tokenizer, the LiteRT model (.tflite), and additional metadata needed to run end-to-end inference with the MediaPipe LLM Inference API.\n",
        "\n",
        "To use the bundler, install the MediaPipe PyPI package (>0.10.14) in this step as it comes with its own set of dependencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rrss72uJ-obD"
      },
      "outputs": [],
      "source": [
        "%pip install mediapipe"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "essOk3jeDmOV"
      },
      "source": [
        "Do a fresh install of `protobuf` and `tensorflow` and restart the Colab runtime to get the latest versions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1LhWaLPOiUTy"
      },
      "outputs": [],
      "source": [
        "%pip uninstall protobuf -y && pip install protobuf\n",
        "%pip uninstall tensorflow -y -q && pip install tensorflow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yHC7dEhR7gFJ"
      },
      "source": [
        "Now, you'll configure and create the Task bundle:\n",
        "\n",
        "1. Update `tflite_model` to point to the newly converted .tflite model in your Colab files\n",
        "2. Update `tokenizer_model` to point to the tokenizer.model that was downloaded from Hugging Face Hub.\n",
        "3. Name your .task file in the `output_filename`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9OxNmqgF-VO0"
      },
      "outputs": [],
      "source": [
        "from mediapipe.tasks.python.genai import bundler\n",
        "\n",
        "config = bundler.BundleConfig(\n",
        "    tflite_model=\"/content/myemoji-gemma-3-270m-it_q8_ekv512.tflite\",     # Point to your converted .tflite model\n",
        "    tokenizer_model=\"/content/myemoji-gemma-3-270m-it/tokenizer.model\",   # Point to the downloaded model's tokenizer.model file\n",
        "    start_token=\"<bos>\",\n",
        "    stop_tokens=[\"<eos>\", \"<end_of_turn>\"],\n",
        "    output_filename=\"/content/myemoji-gemma-3-270m-it.task\",              # Specify the final model filename\n",
        "    prompt_prefix=\"<start_of_turn>user\\n\",\n",
        "    prompt_suffix=\"<end_of_turn>\\n<start_of_turn>model\\n\",\n",
        ")\n",
        "bundler.create_bundle(config)\n",
        "\n",
        "print(f\"Model .task bundle saved to {config.output_filename}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nXtdT_2E5vrO"
      },
      "source": [
        "## Download & run your model on-device\n",
        "\n",
        "Your model is now ready for on-device inference using the MediaPipe LLM Inference API!\n",
        "\n",
        "Download the .task file from your Colab environment to use it in your projects."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NdfSriJW5vNk"
      },
      "outputs": [],
      "source": [
        "from google.colab import files\n",
        "\n",
        "files.download(config.output_filename)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zh1fQTQrcpUh"
      },
      "source": [
        "Try it in the [emoji generation web app](https://github.com/google-gemini/gemma-cookbook/tree/main/Demos/Emoji-Gemma-on-Web/app-mediapipe) which runs the model directly in the browser. You can also explore [documentation](https://ai.google.dev/edge/mediapipe/solutions/genai/llm_inference) for building cross-platform mobile and web apps."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "Convert_Gemma_3_270M_to_LiteRT_for_MediaPipe_LLM_Inference_API.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
